"""--------------------------------------------------------------------------------------------------
Filename:   database.py
Author:     Chris Chester, adapted from Ben Han
Date:       12-09-2013
Description:
    Contains classes to process and store data for the biofeedback application.
    
--------------------------------------------------------------------------------------------------"""

from __future__ import with_statement   # to use the automatic file.close() functionality
import numpy as np
import wx, os, sys, traceback
import time
from config import ConfigData
import pickle


#===============================================================================
# CLASS: Patient
#    This class stores data for each patient. 
#    It is the top level object which contains all other data objects.
#    Must be stored as a native python object. The swig wrappers for the C++ 
#    wx code don't like to be pickled.
#===============================================================================
class Patient:
    def __init__(self, id=False, name='', dob='', gender='', notes='', sessions=[],calValue = None):
        '''
       
        ==============  =====================================
        **Arguments:**
        id              (int) ID number of the patient
        name            (string) Name of the patient
        dob             (string) Date of birth of patient
        gender          (string) Gender of patient
        notes           (string) Notes for patient
        sessions        (array) Array of sessions
        calValue        (int) Calibration value for patient
        ==============  =====================================
        '''
        
        # Initial attributes
        if not id: 
            id = str(int(time.time()))  # Create an arbitrary unique ID if none provided.
            id = id[-8:]                # Uses the number of seconds since 00:00:00 1/1/1970
        self.id = id                        # patient ID number    
        self.name = name                    # patient name
        self.dob = dob                      # patient date of birth
        self.gender = gender                # patient gender
        self.notes = notes                  # storage for general notes, condition, treatment, etc
        self.session_record = sessions      # array containing references to the sessions completed by this patient
        temp = wx.DateTime.Now()
        self.date = temp.Format()           # date created, set to current time & date.
        self.calValue = calValue
            # Format() returns a string in a format that can be parsed
            # back into a wx.DateTime object after unpickling.
        self.path = ""
        self.cfg = ConfigData()
    
    def __str__(self):
        s = 'ID:\t'+ str(self.id)+ '\n'+\
            'Name:\t'+ self.name+ '\n'+\
            'DoB:\t'+ self.dob+ '\n'+\
            'Gender:\t'+ self.gender+ '\n'+\
            'Date:\t'+ self.date+ '\n'+\
            'No. of sessions: '+ str(len(self.session_record))+ '\n'+\
            'Notes:\n'+ self.notes
        return s
       
    def get_calValue(self):
        ''' Receive the stored calibration value
        ==============  =====================================
        '''   
        return self.calValue
    
    def set_calValue(self, callibration):
        ''' Add a session to the internal record
        ==============  =====================================
        **Arguments:**
        callibration       (int) Calibration value to be used
        ==============  =====================================
        '''
        self.calValue = callibration
        
    
        
    def add_session(self, session_rec):
        ''' Add a session to the internal record
        ==============  =====================================
        **Arguments:**
        session_rec     (SessionRecord) Session to be added
        ==============  =====================================
        '''
        self.session_record.append(session_rec)
        return len(self.session_record)
        
    def export(self, parent):
        ''' Opens up a directory selection dialog for user to choose location for output
            CSV files.

        ==============  =====================================
        **Arguments:**
        parent         (wx.frame) Parent window for wx applications
        ==============  =====================================
        '''
        patdir = self.path[:self.path.rfind('\\') + 1]
        dlg = wx.DirDialog(parent,
                           message="Choose a location for output files",
                           defaultPath=patdir)

        path = ""
        if dlg.ShowModal() == wx.ID_OK:
            path = dlg.GetPath()    # directory name for output
            patfilename = self.path[self.path.rfind('\\')+1:self.path.rfind('.pat')]
            outdir = path + "\\" + patfilename + "_output"
            # Create a folder for output
            try:
                os.makedirs(outdir)
            except OSError:
                pass    # directory already exists
            except:
                sys.stderr.write(traceback.format_exc())
            else:
                path = outdir
            
            for i in range(len(self.session_record)):
                sr = self.session_record[i]
                j = self.path.rfind('\\')
                spath = self.path[:j+1] + sr.path
                try:
                    file = open(spath, 'rb')
                    file.flush()
                    s = pickle.load(file)
                    file.close()
                except:
                    msg = ">> Error loading session from file.\n"
                    sys.stderr.write(msg)
                    wx.MessageDialog(None, msg, "Session load error", wx.ICON_WARNING | wx.OK).ShowModal()
                    sys.stderr.write(traceback.format_exc())
                else:
                    outpath = path+ "\\"+ patfilename+ "_Sess"+ str(i+1).rjust(3,'0')
                    s.export(parent, outpath)
        return path
        


#===============================================================================
# CLASS: SessionRecord
#    Database object used to represent session objects in the patient object
#    Actual session objects should be saved in the same folder as the patient
#===============================================================================
class SessionRecord:
    def __init__(self, s_obj):
        '''
        ==============  =====================================
        **Arguments:**
        s_obj           (database.Session) session to extract
                        data from
        ==============  =====================================
        '''

        self.date = s_obj.date
        self.type = s_obj.type
        self.stats = s_obj.stats
        self.path = ""    # path to the file containing the pickled session object
        

#===============================================================================
#  CLASS: Session
#     Stores the session information and configuration file used 
#===============================================================================
class Session:
    def __init__(self, config, pid, type, blocks, trials, stats):
        '''
        ==============  =====================================
        **Arguments:**
        config          (ConfigData) ConfigData object 
                        used for the session
        pid             (int) ID of patient
        type            (string) type of training
        blocks          (list)(Block) list of block objects 
                        for the session
        trials          (list)(Trial) list of trial objects 
                        for the session   
        stats           (Stats) session statistics
                        
        ==============  =====================================
        '''
        
        self.timestamp = str(int(time.time()))
        temp = wx.DateTime.Now()
        self.date = temp.Format()           # time and date of session creation
                            # Date must be stored as a string here, as in Patient class (to be pickled)
        self.config = config                # ConfigData object used for the session
        self.patient = pid                  # ID of patient
        self.type = type                    # type of training (string)
        self.blocks = blocks                # list of block objects for the session
        self.trials = trials                # list of trial objects for the session
        self.stats = stats                  # session statistics
        self.calib_mean = 0                 # mean value of calibration peaks
        
        # Initialise stats object
        self.calc()
        
    def __str__(self):
        s = 'Patient ID,'+ str(self.patient)+ '\n'+\
            'Date,'+ self.date+ '\n'+\
            'Training,'+ self.type+ '\n'+\
            'STATISTICS,'+ '\n'+\
            self.stats.get_col_labels()+ str(self.stats)
        return s
        
    def add_block(self, block):
        '''
        ==============  =====================================
        **Arguments:**
        block           (Block) block to add             
        ==============  =====================================
        '''
        self.blocks.append(block)
        
    def add_trial(self, trial):
        '''
        ==============  =====================================
        **Arguments:**
        trial           (Trial) trial to add             
        ==============  =====================================
        '''
        
        self.trials.append(trial)
        self.calc()
        
    def calc(self):
        '''
        Calculate session statistics.
        The results are written directly to the local Stats object (stats)
        '''
        if not self.stats:
            self.stats = Stats()
        s = self.stats
        s.nhits = 0
        s.ntrials = len(self.trials)
        if s.ntrials == 0:
            pass
        else:
            peaks = []
            for t in self.trials:
                peaks.append(t.peak[1])     # extract the EMG value (float64)
                if t.result:
                    s.nhits+= 1
                
            data = np.array(peaks)
            s.min = data.min()
            s.max = data.max()
            s.mean = data.mean()
            s.hitrate = (s.nhits * 100) / s.ntrials
        
    def export(self, parent, filename=False):
        """ Write session to CSV files. Two files will be created: 
            1- Store stats for the session and each block. 
                It will also record the history of hits and misses. 
            2- The second file will store all the raw EMG data gathered 
                during the session.
                
        ==============  =====================================
        **Arguments:**
        filename        (string) Filename to export data to            
        ==============  =====================================
                
        """
        
        if not filename:
            filename = 'ID_'+ str(self.patient)+ '_'+ self.timestamp
            file_choices = "Comma Seperated Values (*.csv)|*.csv"
            dlg = wx.FileDialog(
                parent,
                message="Export session data to CSV files",
                defaultDir=os.getcwd(),
                defaultFile=filename,
                wildcard=file_choices,
                style=wx.SAVE)
    
            if dlg.ShowModal() == wx.ID_OK:
                filename = dlg.GetPath()
            else:
                filename = ""       # cancelled
                return filename
                
        if filename[-4:] == '.csv':
            filename = filename[:-4]
        self.write_stats(filename)
        self.write_emg(filename)
        return filename
    
    def write_stats(self, filename):
        '''
        write session statistics to a CSV file 
        ==============  =====================================
        **Arguments:**
        filename        (string) Filename to export data to            
        ==============  =====================================
                
        '''
        
        name = filename+ '_stats.csv'
        file = open(name, 'w')
        file.write('<Output for session statistics>\n')
        file.write('Patient ID,'+ str(self.patient)+ '\n')
        file.write('Session date,'+ str(self.date)+ '\n')
        file.write('Training type,'+ str(self.type)+ '\n')
        file.write('PROTOCOLS:\n')
        file.write('Fading,%s\n' % str(self.config.use_fading))
        file.write('Calibration mean peak value,%s\n\n' % str(self.calib_mean))
        file.write('Block Stats\n')
        file.write(self.stats.get_col_labels())
        bi = 1
        for b in self.blocks:       # for each block in the blocks list...
            # write block statistics to file 
            file.write(str(bi)+ ','+ str(b.stats))
            bi += 1
        file.write("\n")
        file.write('Session total,'+ str(self.stats))     # write total session stats
        file.write("\n")
        file.write('Summary of trials\n')
        if len(self.trials)>0:
            file.write(self.trials[0].get_col_labels())
            ti = 1
            for tr in self.trials:
                # write CSV string representation of the history to file
                file.write(str(ti)+ ','+ str(tr))
                ti += 1
        else:
            file.write('(no record of trials)')
        file.flush()      # flush file buffer
        os.fsync(file)
        file.close()
        
    def write_emg(self, filename):
        '''
        write EMG data to a CSV file
        ==============  =====================================
        **Arguments:**
        filename        (string) Filename to export data to            
        ==============  =====================================
                
        '''

        fname = filename+ '_emg.csv'
        file = open(fname, 'w')
        file.write('<Output for session data (raw EMG values)>\n')
        file.write('Patient ID,'+ str(self.patient)+ '\n')
        file.write('Session date,'+ str(self.date)+ '\n')
        file.write('Training type,'+ str(self.type)+ '\n')
        file.write('\n')
        
        # Create a numpy array (2-D) containing all the emg data from trials
        # Every column represents one trial
        
        ntrials = len(self.trials)          # number of trials in the list
        
        nrow = 0                            # Number or rows needed
        
        # Find the maximum number of rows needed

        for i in range(ntrials):

            if len(self.trials[i].emgdata) > nrow:
                nrow =len(self.trials[i].emgdata)
                
        # Compensate for any unfinished data points
        nrow = nrow + 1   
                         
        data_pieces = 3                     # Types of data in each trial (time,emg,plot state)
        
        # initialise a NumPy array filled with zeroes   
        sarray = np.zeros( (nrow, data_pieces*ntrials) )   
        
        # loop through each trial and add its data to the NumPy array
        for c in range(ntrials):              
            emg = self.trials[c].emgdata  
            print "Trial Number"
            print c
            print "Emg Points"
            print len(emg)  
            for r in range(len(emg)):   
                sarray[r,data_pieces*c] = emg[r][0]       # place time value in the array
                sarray[r,data_pieces*c+1] = emg[r][1]       # place emg value in the array3
                
                # if plot state is on for each point, record that
                if emg[r][2]:
                    sarray[r,data_pieces*c+2] = 1     # place plot state in the array
                else:
                    sarray[r,data_pieces*c+2] = 0
                    
        # Write column labels to file stream
        labels = ""
        for c in range(ntrials):
            labels+= 'Time' +','+ 'Trial '+ str(c+1)+ ',Plot visible'
        labels+= '\n'
        file.write(labels)
        
        # Loop over the array and write EMG values to file
        # Each row represents the EMG value for the same time co-ordinate from each trial.
        # The result is a CSV output with one trial per column.
                
        for r in range(nrow):      
            entry = ""
            for c in range(ntrials):   # loop over number of columns
                currentPos = data_pieces*c 
                entry+= "%.2f,%.2f,%s," % ( float(sarray[r,currentPos]),float(sarray[r,currentPos+1]), str(sarray[r,currentPos+2]) )    # format EMG value
            entry+= '\n'
            file.write(entry)
            
        file.flush()    # flush file buffer
        os.fsync(file)  # ensure the file buffers are written to disk
        file.close()


#===============================================================================
#  CLASS: Block
#     Stores block information and raw EMG data 
#===============================================================================
class Block:
    def __init__(self, history=[], stats=False):
        '''
         ==============  =====================================
        **Arguments:**
        history          (list)( (float X, float Y) PEAK, 
                          string TARGET_DETAILS, bool HIT)
                        record of peaks, targets, hits
                        
        stats            (Stats) block statistics
        ==============  =====================================
        
        '''
        
        self.history = history              # record of peaks, targets, hits
        self.stats = stats                  # block statistics
        
    def calc(self):
        '''
        # Calculate block stats
        '''
        if not self.stats:
            self.stats = Stats()
        s = self.stats
        s.nhits = 0
        s.ntrials = len(self.history)
        if s.ntrials == 0:
            pass
        else:
            y = []
            for i in self.history:
                y.append(i[0][1])
                if i[2]: s.nhits+= 1
            data = np.array(y)
            s.min = data.min()
            s.max = data.max()
            s.mean = data.mean()
            s.hitrate = (s.nhits * 100) / s.ntrials


#===============================================================================
# CLASS: Trial
#    Class for storing trial details
#===============================================================================
class Trial:
    def __init__(self, len=None, typ=None, peak=None, target_pos=None,target_size=None, result=None, emgdata=[], y_max=0.0, aspect=True):
        '''
         ==============  =====================================
        **Arguments:**
        len              (int) length of trial in seconds
        typ              (string) training type
        peak             (tuple) a tuple containing the peak 
                         coordinates 
        target_pos       (float,float) X and y positions of the bottom
                         left corner - (xPos,yPos)
        target_size      (float,float) width and height of the 
                         target - (width,height) 
        stats            (Stats) block statistics
        ==============  =====================================
        
        '''
        
        
        self.length = len           # length of trial in seconds
        self.type = typ             # training type (string)
        self.peak = peak            # a tuple containing the peak coordinates
        self.result = result        # set to True if target was hit (boolean)
        self.emgdata = emgdata      # list of tuples storing the raw emg data collected (time, value)
        self.target = self.get_target_tuple(target_pos,target_size)    # tuple storing parameters of a line or rect
        self.y_max = y_max          # float value of y axis range (in uV)
        self.aspect = aspect
        
    def __str__(self):
        '''
        Returns a CSV formatted summary of this trial
        '''
        len = "%i," % int(self.length)
        # get the string representation for the peak
        peak = "%.2f,%.2f," % ( float(self.peak[0]), float(self.peak[1]) )
        # get the string representation for the target
        targ = self.get_target_str()
        # get the string representation for the result
        if self.result:    # if target was hit...
            res = '1,'
        else:
            res = '0,'

        
        # add it all together and return
        s = len+ res+ peak+ targ+ '\n'
        return s
            
    def get_target_tuple(self, pos=[0,0], size=[0,0]):
        '''
        convert the target size and position into a primitive python tuple
         ==============  =====================================
        **Arguments:**
        pos              (float,float) X and y positions of the bottom
                         left corner
        size             (float,float) width and height of the 
                         target
        ==============  =====================================
        
        '''
        xPos = 0
        yPos = 1
        
        if size[yPos] > 0: #if a box, rather than a line
            # for skill (x, y, width, height) where x and y refer to the coordinate of the lower left corner
            return (pos[xPos], pos[yPos], size[xPos], size[yPos])
        
        elif size[xPos] > 0: # if there is a line, rather than a point
            return (pos[yPos], pos[yPos]+size[yPos]) 
        else:
            return ()
        
    def get_target_str(self, output=0):
        '''
        Converts the target tuple into a string
        ==============  =====================================
        **Arguments:**
        output          (int) 0 for CSV formatted output,
                        1 for list format (used in trial
                         details of session pane)
        ==============  =====================================
        '''
        if not self.target:
            return '(none)'
        
        r = self.target
        if len(r) == 2:
            ts = "%.2f" % float(r[0])
            return ts
        elif len(r) == 4:
            cx = r[0]+ r[2]/2.0     # calculate the centre coordinates
            cy = r[1]+ r[3]/2.0     # calculate the centre coordinates
            x = "%.2f" % float(cx)          # centre X
            y = "%.2f" % float(cy)          # centre Y
            w = "%.2f" % float(r[2])        # width
            h = "%.2f" % float(r[3])        # height
            a = "%.2f" % float(r[2]*r[3])   # area
            f = str(self.aspect)            # fixed aspect ratio, yes or no
            pd = "%.2f" % (float(r[0]) + float(r[2]) - float(self.peak[0]))
            
            if output == 0:
                # Convert the matplotlib.patches.Rectangle argument into a string in CSV format
                s = w+ ','+ h+ ','+ a+ ','+  x+ ','+ y+ ','+ pd+ ','+ f+ ','
                return s
            elif output == 1:
                # Convert rectangle to human readable format
                s = 'W:%s, H:%s, Area:%s, Centre:(%s, %s), Fixed aspect ratio: %s' % (w, h, a, x, y, f)
                return s
            else:
                return 'invalid output type'
        else:
            return 'invalid target'

    def emg_str(self):
        '''
        Returns CSV formatted string of the emg data for this trial
        '''
        s = ''
        for e in self.emgdata:
            s+= ("%.2f" % float(e[0])+ ','+ "%.2f" % float(e[1])+ '\n')
        return s
        
    def get_col_labels(self):
        '''
        Returns labels for the columns
        '''
        s = ''
        if len(self.target) == 2:
            s = 'Trial Number,Length(secs),Result(1=hit),Peak X (secs),Peak Y (uV),Plot size X (pixels),Plot size Y (pixels),Y axis max (uV),Target amplitude\n'
        elif len(self.target) == 4:
            s = 'Trial Number,Length(secs),Result(1=hit),Peak X (secs),Peak Y (uV),Plot size X (pixels),Plot size Y (pixels),Y axis max (uV),Target W (secs),Target H (uV),Target Area,Centre X (secs),Centre Y (uV),Peak delay,Fixed Aspect Ratio\n'
        return s
    
    def get_summary(self):
        '''
        Prints a summary of the trial parameter to sys.stdout
        '''
        trial_s = '> Trial summary\n'
        trial_s += 'Training mode: %s\n' % str(self.type)
        trial_s += 'Length: %.2f seconds\n' % float(self.length)
        trial_s += 'Y axis max: %.2f uV\n' % self.y_max
        trial_s += 'Target: %s\n' % self.get_target_str(1)
        trial_s += 'Peak: (%.2f, %.2f)\n' % (float(self.peak[0]), float(self.peak[1]))
        trial_s += 'Result: %s\n' % str(self.result)
        return trial_s
    


#===============================================================================
#  CLASS: Stats
#     Used by Block and Session to store basic stats 
#===============================================================================
class Stats:
    def __init__(self):
        self.ntrials = 0
        self.min = 0
        self.max = 0
        self.mean = 0
        self.nhits = 0
        self.hitrate = 0
        
    def __str__(self):
        '''
        The return value is a string containing the class attributes in CSV format.
        '''
        s = "%i" % int(self.ntrials)+ ','+ "%.2f" % float(self.min)+ ','+ "%.2f" % float(self.max)+\
            ','+ "%.2f" % float(self.mean)+ ','+ "%i" % int(self.nhits)+ ','+ "%.2f" % float(self.hitrate)+ '\n'
        return s
    
    def __repr__(self):
        '''
        The representation entered into the patient database
        '''
        s = ''
        s+= 'No. of trials: '+ "%i" % int(self.ntrials)+ '\n'
        s+= 'Min peak:      '+ "%.2f" % float(self.min)+ '\n'
        s+= 'Max peak:      '+ "%.2f" % float(self.max)+ '\n'
        s+= 'Mean peak:     '+ "%.2f" % float(self.mean)+ '\n'
        s+= 'No. of hits:   '+ "%i" % int(self.nhits)+ '\n'
        s+= 'Hit rate:      '+ "%.2f" % float(self.hitrate)+ '\n'
        return s
    
    def get_col_labels(self):
        s = 'Block No.,No. trials,Min peak,Max peak,Mean peak,No. hits,Hitrate\n'
        return s

# End of database.py
